import math
import sys
import warnings

from matplotlib import pyplot
from matplotlib.colors import LogNorm
import numpy

from h1d import Histogram1D
from h2d import Histogram2D



class Histogram3D(object):
    """
    Three-dimensional histogram object.
    
    Parameters
    ----------
    name   : string
        A name for the histogram
    nbins  : sequence of length 3
        (``number of bins in x``, ``number of bins in y``, ``number of bins in z``)
    arange : a sequence of two sequences of length 3
        ((``x axis minimum``, ``x axis maximum``), (``y axis minimum``, ``y axis maximum``), (``z axis minimum, z axis maximum``))
    xlabel : string
        Label for the x-axis
    ylabel : string
        Label for the y-axis
    zlabel : string
        Label for the z-axis
    xunits : string 
        Units for the x-axis
    yunits : string 
        Units for the y-axis
    zunits : string 
        Units for the z-axis


    Optional Parameters
    -------------------
    title : string 
        Title of the histogram (default is "`xlabel` vs `ylabel`")
    fontsize  :  float | {'xx-small' , 'x-small' , 'small' , 'medium' , 'large' , 'x-large' , 'xx-large'}
        Size of text when drawn
    ext       : {'png' , 'gif' , 'jpg'}
        Default file extension to add to name when histogram is drawn and saved to file
    bufsize   : int
            length of internal data buffer; larger bufsizes are faster for large datasets but take up more memory
    assoc_cols: 3-tuple of string
        Three column names (x, y, z) which are associated with this histogram
    selection : string
        A PyTables-style selection which is associated with this histogram
    """

    def __init__(self, name, nbins, aranges, 
                 xlabel, ylabel, zlabel,
                 xunits, yunits, zunits,
                 title='', fontsize='x-large', ext='png', bufsize=1000000, 
                 assoc_cols=[]):
        self.name      = name
        self.nxbins    = nbins[0]
        self.nybins    = nbins[1]
        self.nzbins    = nbins[2]
        self.xrange    = aranges[0]
        self.yrange    = aranges[1]
        self.zrange    = aranges[2]
        self.xlabel    = xlabel
        self.ylabel    = ylabel
        self.zlabel    = zlabel
        self.xunits    = xunits
        self.yunits    = xunits
        self.zunits    = zunits
        self.title     = title
        self.fontsize  = fontsize
        self.ext       = ext
        self.assoc_cols = assoc_cols

        self._i        = 0;
        self._tot_i    = 0;

        self._bufsize  = bufsize
        self.low_edges_x = None
        self.low_edges_y = None
        self.low_edges_z = None
        self.bin_width = (float(self.xrange[1] - self.xrange[0]) / self.nxbins,
                          float(self.yrange[1] - self.yrange[0]) / self.nybins,
                          float(self.zrange[1] - self.zrange[0]) / self.nzbins)
                          

        self._bins  = None
        self._buf   = None
        self._wbuf  = None
        self._sum   = 0

        self._xmin    = 1e250
        self._xmax    = -1e250
        self._ymin    = 1e250
        self._ymax    = -1e250
        self._zmin    = 1e250
        self._zmax    = -1e250
        self._mean   = 0.0 
        self._var    = 0.0

    def _setup_arrays(self):
        if self._bins is None:
            self._bins  = numpy.zeros((self.nxbins, self.nybins, self.nzbins))

        if self._buf is None:
            self._buf   = numpy.zeros((self._bufsize, 3))

        if self._wbuf is None:
            self._wbuf  = numpy.ones((self._bufsize,))

        if self.low_edges_x is None:
            self.low_edges_x, bin_width_x = numpy.linspace(self.xrange[0], self.xrange[1], self.nxbins, retstep=True)

        if self.low_edges_y is None:
            self.low_edges_y, bin_width_y = numpy.linspace(self.yrange[0], self.yrange[1], self.nybins, retstep=True)

        if self.low_edges_z is None:
            self.low_edges_z, bin_width_z = numpy.linspace(self.zrange[0], self.zrange[1], self.nzbins, retstep=True)
               

    def draw(self, hold=False, same=False, cmap='hot', **kwargs):
        """
        Draw this histogram. 

        ``Histogram3D`` is meant primarily for slicing and projecting, so drawing is not implemented here.
        """
        warnings.warn("Attempt to draw 3D histogram `{0}'--ignoring...".format(self.name))
        #raise NotImplementedError("Drawing of 3D histograms is not available.")
        
    
    def fill(self, value, weight=1, loud=False):
        """
        Fill a histogram with values.
                
        Parameters
        ----------
        
        value : ``(x, y, z)`` | sequence of ``(x, y, z)``
            Values to be binned

        Optional Parameters
        -------------------
        weight : int | sequence of int
            Weights to apply when binning values
        loud   : bool
            Print a progress indicator when binning multiple values
        """

        if self._buf is None:
            self._setup_arrays()

        try:
            self._buf[self._i] = value
        except IndexError:
            self.flush()
            try:
                self._buf[self._i] = value
            except ValueError:
                return self._fill_array(value, weight, loud)
        except ValueError:
            return self._fill_array(value, weight, loud)

        self._wbuf[self._i] = weight
        self._i += 1

    def _fill_array(self, values, weights, loud):
        if not len(values):
            return

        valshape = numpy.shape(values)

        if len(valshape) != 2 or 3 not in valshape:
            msg = "Trying to fill 3D histogram with array of shape `{0}'; valid shapes are (2,n) or (n,2)"
            raise ValueError(msg.format(valshape))
        
        if valshape[1] != 3:
            try:
                values = values.T
            except AttributeError:
                values = numpy.asarray(values).T

        try:
            x = iter(weights)
        except TypeError:
            pass
        else:
            return self._fill_array_weights(values, weights, loud)
        
        self.flush()
        
        val_len = len(values)

        if val_len < self._bufsize:
            self._i = val_len
            self._buf[:self._i] = values[:self._i]
            return

        for i in xrange(0, val_len, self._bufsize):
            if loud:
                sys.stderr.write("Histogram3D `{0}': {1}/{2} events binned\r".format(self.name, i, val_len))
                sys.stderr.flush()
            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            self.flush()

    def _fill_array_weights(self, values, weights, loud):
        self.flush()
        
        self._i = len(values)
        if self._i < self._bufsize:
            self._i = len(values)
            self._buf[:self._i]  = values
            self._wbuf[:self._i] = weights
            return

        for i in xrange(0, val_len, self._bufsize):
            if loud:
                sys.stderr.write("Histogram3D `{0}': {1}/{2} events binned\r".format(self.name, i, val_len))
                sys.stderr.flush()

            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            self._wbuf[:self._i] = weights[i : i + self._bufsize]
            self.flush()

    def clear(self):
        """Set all bin contents to zero."""        

        if self._buf is not None:
            self._buf   = numpy.zeros((self._bufsize,))

        if self._bins is not None:
            self._bins  = numpy.zeros((self.nxbins,self.nybins))

        self._mean   = 0.0 
        self._var    = 0.0
        self._xmin    = 1e250
        self._xmax    = -1e250
        self._ymin    = 1e250
        self._ymax    = -1e250
        self._zmin    = 1e250
        self._zmax    = -1e250
        self._i      = 0
        self._tot_i  = 0

            
    def flush(self):
        """
        Bin events in buffer and clear the buffer.
        """

        if self._buf is None or self._bins is None:
            self._setup_arrays()

        if self._i == 0:
            return
        
        buf_to_flush = self._buf[:self._i]

        bufxmax = max(buf_to_flush.T[0])
        if bufxmax > self._xmax:
            self._xmax = bufxmax
        bufxmin = min(buf_to_flush.T[0])
        if bufxmin < self._xmin:
            self._xmin = bufxmin

        bufymax = max(buf_to_flush.T[1])
        if bufymax > self._ymax:
            self._ymax = bufymax
        bufymin = min(buf_to_flush.T[1])
        if bufymin < self._ymin:
            self._ymin = bufymin

        bufzmax = max(buf_to_flush.T[2])
        if bufzmax > self._zmax:
            self._zmax = bufzmax
        bufzmin = min(buf_to_flush.T[2])
        if bufzmin < self._zmin:
            self._zmin = bufzmin

        
        self._mean = (self._mean * self._tot_i + numpy.mean(buf_to_flush) * self._i) / (self._tot_i + self._i)
        self._var  += sum((buf_to_flush - self._mean)**2)
 
        bufbins, edges = numpy.histogramdd(buf_to_flush,
                                           bins=(self.nxbins, self.nybins, self.nzbins),
                                           weights=self._wbuf[:self._i],
                                           range=(self.xrange, self.yrange, self.zrange))

        
        self._bins += bufbins
        self._tot_i += self._i
        self._i = 0

    def project_to_xy(self, bin_start=0, bin_stop=None):
        """
        Project the histogram onto the x-y plane

        Optional Parameters
        -------------------
        bin_start : int
            z-bin index of lower bound of projection
        bin_stop  : int
            z-bin index of upper bound of projection

        Returns
        -------
        a Histogram2D object.
        """

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        if bin_stop is None:
            bin_stop = self._bins.shape[2]

        binst = numpy.transpose(self._bins, axes=(2,1,0))[bin_start:bin_stop]
        bins  = numpy.sum(binst, 0)
        return Histogram2D(self.name + "_pxy", bins.shape, (self.xrange, self.yrange), 
                           self.xlabel, self.ylabel, self.xunits, self.yunits,
                           self.xlabel + ' vs ' + self.ylabel, bins=bins.T)# <<<<<<<<<<<<<<<<<<<<<,)

    def project_to_xz(self, bin_start=0, bin_stop=None):
        """
        Project the histogram onto the x-z plane

        Optional Parameters
        -------------------
        bin_start : int
            y-bin index of lower bound of projection
        bin_stop  : int
            y-bin index of upper bound of projection

        Returns
        -------
        a Histogram2D object.
        """    

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        if bin_stop is None:
            bin_stop = self._bins.shape[1]

        binst = numpy.transpose(self._bins, axes=(1,2,0))[bin_start:bin_stop]
        bins  = numpy.sum(binst, 0)
        return Histogram2D(self.name + "_pxz", bins.shape, (self.xrange, self.zrange), 
                           self.xlabel, self.zlabel, self.xunits, self.zunits,
                           self.xlabel + ' vs ' + self.zlabel, bins=bins.T)

    def project_to_yz(self, bin_start=0, bin_stop=None):
        """
        Project the histogram onto the x-z plane

        Optional Parameters
        -------------------
        bin_start : int
            x-bin index of lower bound of projection
        bin_stop  : int
            x-bin index of upper bound of projection

        Returns
        -------
        a Histogram2D object.
        """

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()
        
        if bin_stop is None:
            bin_stop = self._bins.shape[0]

        binst = numpy.transpose(self._bins, axes=(0,2,1))[bin_start:bin_stop]
        bins  = numpy.sum(binst, 0)
        return Histogram2D(self.name + "_pyz", bins.shape, (self.yrange, self.zrange), 
                           self.ylabel, self.zlabel, self.yunits, self.zunits,
                           self.ylabel + ' vs ' + self.zlabel, bins=bins.T)
        
    
    def set_buffer_size(self, bufsize):
        """
        Set internal buffer size to new value.
        
        Parameters
        ----------
        
        bufsize : int
            desired new buffer size
        """

        self._bufsize = bufsize

        if self._buf is not None:            
            self.flush()
            
            self._buf     = numpy.zeros((self._bufsize,))
            self._wbuf    = numpy.ones((self._bufsize,))

        
    def __getstate__(self):
        self.flush()

        outdict = {}
        outdict.update(self.__dict__.items())

        outdict['low_edges_x'] = None
        outdict['low_edges_y'] = None
        outdict['low_edges_z'] = None
        outdict['_bins'] = None
        outdict['_buf'] = None
        outdict['_wbuf'] = None
        
        return outdict


    def bins(self):
        """Returns a reference to the array of bin contents."""

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        return self._bins


    def max_bin_index(self):
        """
        Returns the indices ``(ix, iy)`` of the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()

        return numpy.unravel_index(numpy.argmax(self._bins, self._bins.shape))
                
    def max_bin_content(self):
        """
        Returns the number of counts in the bin with the most counts.
        """
        assert self._bins is not None, "Histogram is empty"

        self.flush()

        i, j, k = self.max_bin_index()
        return self._bins[i][j][k]

    def max_bin_xy(self):
        """
        Returns the (x, y, z) coordinates of the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"        

        self.flush()

        return self.bin_center(self.max_bin_index())

    def bin_center(self, ix, iy, iz):
        """
        Returns the x-y coordinates ``(x, y)`` of the center of the bin at ``(ix, iy)``.
        """

        if self.low_edges_x is None or self.low_edges_y is None or self.low_edges_z is None:
            self._setup_arrays()

        return (self.low_edges_x[ix] + self.bin_width[0] / 2.0, 
                self.low_edges_y[iy] + self.bin_width[1] / 2.0,
                self.low_edges_z[iz] + self.bin_width[2] / 2.0)

    def bin_index(self, x, y, z):
        """
        Return the indices ``(ix, iy)`` of the bin that ``(x, y)`` would be placed into.
        """ 

        if self.low_edges_x is None or self.low_edges_y is None or self.low_edges_z is None:
            self._setup_arrays()
       
        return (bisect.bisect(self.low_edges_x, x), 
                bisect.bisect(self.low_edges_y, y), 
                bisect.bisect(self.low_edges_z, z))
    
    def data_xmax(self):
        """
        Returns the current maximum x-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._xmax

    def data_xmin(self):
        """
        Returns the current minimum x-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._xmin

    def data_ymax(self):
        """
        Returns the current maximum y-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._ymax

    def data_ymin(self):
        """
        Returns the current minimum y-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._ymin

    def data_mean(self):
        """
        Returns the mean ``(mu_x, mu_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._mean

    def data_var(self):
        """
        Returns the variance ``(var_x, var_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._var

    def data_std(self):
        """
        Returns the standard deviation ``(std_x, std_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return math.sqrt(self._var / float(self._tot_i))

    def nentries(self):
        """
        Returns the total number of values binned.
        """

        if self._bins is None:
            return 0

        self.flush()
        return self._tot_i

    def __len__(self):
        return self.nentries()


    def __iadd__(self, rhs):     
        if self._bins is None:
            self._setup_arrays()

        if rhs._bins is None:
            rhs._setup_arrays()

        self.flush()
        rhs.flush()

        if len(rhs) == 0:
            return self
        
        tot_entries = self._tot_i + rhs._tot_i
        self._mean = self._mean * (self._tot_i / float(tot_entries)) + rhs._mean * (rhs._tot_i / float(tot_entries))
        self._var  = self._var * (self._tot_i / float(tot_entries)) + rhs._var * (rhs._tot_i / float(tot_entries))

        self._tot_i = tot_entries
        
        if self._xmax < rhs._xmax:
            self._xmax = rhs._xmax

        if self._ymax < rhs._ymax:
            self._ymax = rhs._ymax

        if self._zmax < rhs._zmax:
            self._zmax = rhs._zmax

        if self._xmin < rhs._xmin:
            self._xmin = rhs._xmin

        if self._ymin < rhs._ymin:
            self._ymin = rhs._ymin

        if self._zmin < rhs._zmin:
            self._zmin = rhs._zmin
            
        self._bins += rhs._bins
        
        return self

        
            
